import os
import csv
import pandas as pd
from jieba import analyse


def get_stopwords(file_pth):
    """读取停用词表"""
    with open(file_pth, mode="r", encoding="utf-8") as fp:
        stopwords_list = [i.strip() for i in fp.readlines()]
    return stopwords_list


def insert_dict(d, index, value):
    if index not in d:
        d[index] = []
    d[index].append(value)


def stripdata(data_pth, dict_pth, out_pth):
    # 读取停用词表
    stopwords_list = get_stopwords(os.path.join(dict_pth, "baidu_stopwords.txt"))

    # 读取关键词表
    given_tag_tb = pd.read_csv(os.path.join(data_pth, "Movie_tag.csv"),encoding='utf-8')
    given_tag_tb.set_index("id", inplace=True)

    # 从爬取的数据中，分析关键词，并和给定的关键词合并
    tags_id_dict = dict()  # 倒排表
    id_name_dict = dict()  # id-name表
    id_tags_dict = dict()  # id-tags表
    with open(
        os.path.join(out_pth, "movie.csv"), mode="r", encoding="utf-8"
    ) as fp_movie:
        text_csv = csv.reader(fp_movie)
        next(text_csv)
        for row in text_csv:
            if row:
                _id, _name, _intro = int(row[0]), row[1], row[3]
                # 处理id-name
                id_name_dict[_id] = _name

                # 处理id-tag
                try:  # 读取给定的tag
                    given_tag = given_tag_tb.loc[_id]["tag"].split(",")
                except:
                    given_tag = []
                given_tag = set(given_tag)
                # print(given_tag)
                tag_list = analyse.extract_tags(  # 从简介中提取关键词
                    _intro, topK=5, withWeight=False, allowPOS=()
                )
                # 删除在停用词表中的tag，并去重
                #!注意，此后此后修改id_tags_dict[_id]时，given_tag也会被修改
                id_tags_dict[_id] = given_tag
                temp = []
                for tag in tag_list:
                    if tag not in stopwords_list:
                        temp.append(tag)
                id_tags_dict[_id].update(temp)

                # 更新倒排表
                for tag in temp:
                    insert_dict(tags_id_dict, tag, str(_id))

    # 写入id-name表
    id_name_tb = pd.DataFrame(
        {"id": id_name_dict.keys(), "name": id_name_dict.values()}
    )
    id_name_tb.set_index("id", inplace=True)
    id_name_tb.to_csv(os.path.join(out_pth, "id_name.csv"))

    # 写入id-tag表
    id_tags_tb = pd.DataFrame(
        {
            "id": id_tags_dict.keys(),
            "tag": [", ".join(tag) for tag in id_tags_dict.values()],
        }
    )
    id_tags_tb.set_index("id", inplace=True)
    id_tags_tb.to_csv(os.path.join(out_pth, "Movie_tag_add.csv"))

    # 写入倒排表
    tags_id_tb = pd.DataFrame(
        {
            "tag": tags_id_dict.keys(),
            "id": [", ".join(sorted(_id)) for _id in tags_id_dict.values()],
        }
    )
    tags_id_tb.set_index("tag", inplace=True)
    tags_id_tb.to_csv(os.path.join(out_pth, "inverted_index.csv"))


if __name__ == "__main__":
    pth = os.path.split(os.path.realpath(__file__))[0]
    data_pth = os.path.join(pth, "data")
    out_pth = os.path.join(pth, "output")
    dict_pth = os.path.join(pth, "stopwords")

    stripdata(data_pth, dict_pth, out_pth)
